import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional, Dropout
from tensorflow.keras.callbacks import ReduceLROnPlateau
import tushare as ts
import math
import random

# 设置随机种子，保证结果的可复现性
random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

# 设置中文字体为黑体，解决负号显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 获取股票数据的函数，使用Tushare的数据接口
def get_stock_data(token, ts_code, start_date, end_date, fields):
    try:
        pro = ts.pro_api(token)  # 请替换为自己的Tushare token
        data = pro.daily(ts_code=ts_code,
                         start_date=start_date,
                         end_date=end_date,
                         fields=fields)
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None


# 数据预处理函数，进行缺失值处理、数据标准化等操作
def preprocess_data(data):
    df = data.copy()
    # 将交易日期列设置为索引，并转换为日期时间类型
    df.set_index('trade_date', inplace=True)
    df.index = pd.to_datetime(df.index)

    # 去除包含缺失值的行
    df.dropna(inplace=True)

    # 提取需要的列（开盘价、最高价、最低价、收盘价、成交量）
    columns = ['open', 'high', 'low', 'close', 'vol', 'pre_close', 'change', 'pct_chg']
    df = df[columns]

    # 数据标准化
    scaler = StandardScaler()
    if df.shape[0] == 0:
        print("数据为空，无法进行标准化处理！")
        exit(1)

    df_scaled = scaler.fit_transform(df)
    return df, df_scaled, scaler


# 创建训练数据集和测试数据集的函数，按照时间序列划分数据
def create_dataset(dataset, time_steps=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - time_steps - 1):
        a = dataset[i:(i + time_steps), :]
        dataX.append(a)
        dataY.append(dataset[i + time_steps, -1])  # 预测目标是收盘价
    return np.array(dataX), np.array(dataY)


# 构建并训练LSTM模型的函数，支持不同模型类型选择
def build_and_train_model(model_type, X_train, y_train, X_test, y_test):
    model = Sequential()
    if model_type == 1:
        # 单层 LSTM
        model.add(LSTM(units=100, activation='relu',
                       input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(units=1))
    elif model_type == 2:
        # 多层 LSTM
        model.add(LSTM(units=100, activation='relu', return_sequences=True,
                       input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dropout(0.2))  # 加入Dropout层
        model.add(LSTM(units=100, activation='relu'))
        model.add(Dense(1))
    elif model_type == 3:
        # 双向 LSTM
        model.add(Bidirectional(LSTM(100, activation='relu'),
                                input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(1))

    model.compile(loss='mean_squared_error', optimizer='adam')

    # 学习率动态调整
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, min_lr=1e-6)
    history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=100, batch_size=32, callbacks=[reduce_lr], verbose=1)

    return model, history


# 计算均方根误差（RMSE），用于评估模型性能
def calculate_rmse(y_true, y_pred):
    min_len = min(len(y_true), len(y_pred))  # 获取较小的长度
    y_true = y_true[:min_len]  # 修剪 y_true
    y_pred = y_pred[:min_len]  # 修剪 y_pred

    return math.sqrt(((y_true - y_pred) ** 2).mean())


# 绘制损失曲线的函数
def plot_loss_curve(history):
    fig = go.Figure()

    fig.add_trace(go.Scatter(x=list(range(len(history.history['loss']))), 
                             y=history.history['loss'], mode='lines', name='训练损失'))
    fig.add_trace(go.Scatter(x=list(range(len(history.history['val_loss']))), 
                             y=history.history['val_loss'], mode='lines', name='验证损失'))

    fig.update_layout(title="模型训练损失曲线",
                      xaxis_title="Epoch",
                      yaxis_title="损失值")
    fig.show(renderer="browser")  # 强制在浏览器中显示图表


# 绘制股价预测图
def plot_predictions(df, train_size, time_steps, y_test, test_predict):
    fig = go.Figure()

    # 绘制真实值
    fig.add_trace(go.Scatter(x=df.index[train_size + time_steps + 1:], 
                             y=y_test, mode='lines', name='真实值'))
    
    # 绘制预测值
    fig.add_trace(go.Scatter(x=df.index[train_size + time_steps + 1:], 
                             y=test_predict, mode='lines', name='预测值'))

    fig.update_layout(title="真实值与预测值对比",
                      xaxis_title="日期",
                      yaxis_title="收盘价")
    fig.show(renderer="browser")  # 强制在浏览器中显示图表


if __name__ == "__main__":
    # Tushare平台获取的token，需替换为你自己的真实token
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'
    ts_code = "600519.SH"  # 贵州茅台股票代码，可替换为其他感兴趣的股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)

    # 检查获取的数据是否为空
    if stock_data is None or stock_data.empty:
        print("没有获取到有效的股票数据！")
        exit(1)

    # 数据预处理
    df, df_scaled, scaler = preprocess_data(stock_data)

    # 划分训练集和测试集
    print("数据总长度:", len(df_scaled))
    train_size = int(len(df_scaled) * 0.8)
    test_size = len(df_scaled) - train_size
    print("训练集大小:", train_size)
    print("测试集大小:", test_size)
    train_data, test_data = df_scaled[0:train_size, :], df_scaled[train_size:len(df_scaled), :]
    time_steps = 20  # 增加时间步长
    X_train, y_train = create_dataset(train_data, time_steps)
    X_test, y_test = create_dataset(test_data, time_steps)
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], df.shape[1]))
    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], df.shape[1]))

    # 选择模型类型（1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM）
    model_type = 2  # 尝试使用多层 LSTM
    model, history = build_and_train_model(model_type, X_train, y_train, X_test, y_test)

    # 绘制损失曲线
    plot_loss_curve(history)

    # 预测股价
    test_predict = model.predict(X_test)

    # 反标准化预测结果和真实结果
    test_predict = scaler.inverse_transform(np.concatenate((np.zeros((test_predict.shape[0], df.shape[1] - 1)), test_predict), axis=1))[:, -1]
    y_test_actual = scaler.inverse_transform(np.concatenate((np.zeros((y_test.shape[0], df.shape[1] - 1)), y_test.reshape(-1, 1)), axis=1))[:, -1]

    # 计算RMSE
    rmse = calculate_rmse(y_test_actual, test_predict)
    print(f"测试集的RMSE: {rmse}")

    # 绘制股价预测结果与真实股价对比图
    plot_predictions(df, train_size, time_steps, y_test_actual, test_predict)
